#include <stddef.h>
#include <stdint.h>
#include <stdio.h>
#include <string.h>
#include <sys/types.h>
#include <tss2/tss2_rc.h>
#include <tss2/tss2_sys.h>
#include <tss2/tss2_tctildr.h>

#include <openssl/err.h>
#include <openssl/evp.h>
#include <openssl/rand.h>

#ifndef TCTI_NAME
#define TCTI_NAME "mssim"
#endif

TSS2_TCTI_CONTEXT *tcti_context = NULL;
TSS2_SYS_CONTEXT *sys_context = NULL;

// Record handles for cleanup
typedef struct HandleRecorder {
    TPM2_HANDLE handles[128];
    uint16_t count;
} HandleRecorder;

static HandleRecorder handle_recorder = {
    .count = 0, .handles = {TPM2_RH_NULL} // Initialize with NULL handle
};

// Check if the return code indicates success, clean up if not
#define CHECK_PASSED(rc)                                                       \
    {                                                                          \
        if ((rc) != TPM2_RC_SUCCESS) {                                         \
            printf("FAILED! [%s:%u] %s\n", __FUNCTION__, __LINE__,             \
                   Tss2_RC_Decode(rc));                                        \
            cleanup(rc);                                                       \
        } else {                                                               \
            printf("PASSED! [%s:%u]\n", __FUNCTION__, __LINE__);               \
        }                                                                      \
    }

#define CHECK_OSSL(cond, func)                                                 \
    {                                                                          \
        if (!(cond)) {                                                         \
            fprintf(stderr, "ERROR: OpenSSL call %s failed! [%s:%u]\n", func,  \
                    __FUNCTION__, __LINE__);                                   \
            ERR_print_errors_fp(stderr);                                       \
            cleanup(1);                                                        \
        }                                                                      \
    }

// Initialize TCTI Context
TSS2_RC init_tcti_context(TSS2_TCTI_CONTEXT **tcti_context) {
    TSS2_RC rc = Tss2_TctiLdr_Initialize(TCTI_NAME, tcti_context);
    return rc;
}

// Initialize SAPI Context
TSS2_RC init_sys_context(TSS2_TCTI_CONTEXT *tcti_context,
                         TSS2_SYS_CONTEXT **sys_context) {
    if (tcti_context == NULL) {
        return TSS2_SYS_RC_BAD_REFERENCE;
    }

    size_t sys_context_size = Tss2_Sys_GetContextSize(0);
    *sys_context = (TSS2_SYS_CONTEXT *)malloc(sys_context_size);
    if (*sys_context == NULL) {
        return TSS2_BASE_RC_MEMORY;
    }

    TSS2_ABI_VERSION abi_version = TSS2_ABI_VERSION_CURRENT;
    TSS2_RC rc = Tss2_Sys_Initialize(*sys_context, sys_context_size,
                                     tcti_context, &abi_version);
    if (rc != TSS2_RC_SUCCESS) {
        free(*sys_context);
        *sys_context = NULL;
    }

    return rc;
}

// Startup TPM
TSS2_RC startup_tpm(TSS2_SYS_CONTEXT *sys_context) {
    if (sys_context == NULL) {
        return TSS2_SYS_RC_BAD_REFERENCE;
    }

    TSS2_RC rc =
        Tss2_Sys_GetCapability(sys_context, NULL, TPM2_CAP_TPM_PROPERTIES,
                               TPM2_PT_MANUFACTURER, 1, NULL, NULL, NULL);
    if (rc == TPM2_RC_INITIALIZE) {
        rc = Tss2_Sys_Startup(sys_context, TPM2_SU_CLEAR);
    }

    return rc;
}

// Cleanup resources
static void cleanup(int rc) {
    for (uint16_t i = 0; i < handle_recorder.count; i++) {
        if (handle_recorder.handles[i] != TPM2_RH_NULL) {
            TSS2_RC rc_flush =
                Tss2_Sys_FlushContext(sys_context, handle_recorder.handles[i]);
            if (rc_flush != TSS2_RC_SUCCESS) {
                printf("FlushContext Failed for handle 0x%08x\n",
                       handle_recorder.handles[i]);
            }
        }
    }
    handle_recorder.count = 0;
    if (tcti_context != NULL) {
        Tss2_TctiLdr_Finalize(&tcti_context);
        free(tcti_context);
    }
    if (sys_context != NULL) {
        Tss2_Sys_Finalize(sys_context);
        free(sys_context);
    }
    exit(rc != TSS2_RC_SUCCESS);
    return;
}

// Create a RSA 2048 key
TSS2_RC create_key_rsa_2048(TSS2_SYS_CONTEXT *sys_context, const char *password,
                            TPM2_HANDLE *key_handle, TPM2B_PUBLIC *out_public) {
    // CreatePrimary Input Parameters
    TSS2L_SYS_AUTH_COMMAND cmd_auths_array = {
        .count = 1, .auths = {{.sessionHandle = TPM2_RH_PW}}};
    TPM2B_PUBLIC in_public = {
        .size = 0,
        .publicArea =
            {
                .type = TPM2_ALG_RSA,
                .nameAlg = TPM2_ALG_SHA256,
                .objectAttributes = TPMA_OBJECT_DECRYPT |
                                    TPMA_OBJECT_USERWITHAUTH |
                                    TPMA_OBJECT_SENSITIVEDATAORIGIN,
                .parameters.rsaDetail =
                    {
                        .scheme = {.scheme = TPM2_ALG_OAEP,
                                   .details = {.oaep = {.hashAlg =
                                                            TPM2_ALG_SHA256}}},
                        .keyBits = 2048,
                        .exponent = 0,
                        .symmetric = {.algorithm = TPM2_ALG_NULL},
                    },
            },
    };
    UINT16 pwd_size = strlen(password);
    TPM2B_SENSITIVE_CREATE in_sensitive = {
        .sensitive.userAuth =
            {
                .size = pwd_size,
                .buffer = {},
            },
    };
    memcpy(in_sensitive.sensitive.userAuth.buffer, password, pwd_size);
    TPM2B_DATA outside_info = {
        .size = 0,
        .buffer = {},
    };
    TPML_PCR_SELECTION creation_pcr = {
        .count = 0,
    };

    // CreatePrimary Output Parameters
    TPM2B_CREATION_DATA creation_data = {0};
    TPM2B_DIGEST creation_hash = {.size = sizeof(creation_hash.buffer)};
    TPMT_TK_CREATION creation_ticket = {0};
    TPM2B_NAME name = {0};
    TSS2L_SYS_AUTH_RESPONSE rsp_auths_array = {0};

    TSS2_RC rc = Tss2_Sys_CreatePrimary(
        sys_context, TPM2_RH_OWNER, &cmd_auths_array, &in_sensitive, &in_public,
        &outside_info, &creation_pcr, key_handle, out_public, &creation_data,
        &creation_hash, &creation_ticket, &name, &rsp_auths_array);
    return rc;
};

TSS2_RC load_pubkey(TSS2_SYS_CONTEXT *sys_context, TPM2B_PUBLIC *in_public,
                    TPM2_HANDLE *key_handle) {

    TSS2_RC rc = Tss2_Sys_LoadExternal(sys_context, NULL, NULL, in_public,
                                       TPM2_RH_NULL, key_handle, NULL, NULL);
    return rc;
}

// RSA Encrypt data using public key
TSS2_RC rsa_pubkey_encrypt(TSS2_SYS_CONTEXT *sys_context, const char *password,
                           TPM2_HANDLE key_handle,
                           TPM2B_PUBLIC_KEY_RSA *in_data,
                           TPM2B_PUBLIC_KEY_RSA *out_data) {
    // Encrypt Input Parameters
    TPMT_RSA_DECRYPT in_scheme = {
        .scheme = TPM2_ALG_OAEP,
        .details = {.oaep = {.hashAlg = TPM2_ALG_SHA256}}};
    TPM2B_DATA label = {0};
    TSS2_RC rc = Tss2_Sys_RSA_Encrypt(sys_context, key_handle, NULL, in_data,
                                      &in_scheme, &label, out_data, NULL);
    return rc;
}

// RSA Decrypt data using private key
TSS2_RC rsa_prvkey_decrypt(TSS2_SYS_CONTEXT *sys_context, const char *password,
                           TPM2_HANDLE key_handle,
                           TPM2B_PUBLIC_KEY_RSA *in_data,
                           TPM2B_PUBLIC_KEY_RSA *out_data) {
    // EncryptDecrypt Input Parameters
    TSS2L_SYS_AUTH_COMMAND cmd_auths_array = {
        .count = 1,
        .auths[0] = {
            .sessionHandle = TPM2_RH_PW,
            .hmac = {.size = strlen(password)},
        }};
    memcpy(cmd_auths_array.auths[0].hmac.buffer, password, strlen(password));
    TPMT_RSA_DECRYPT in_scheme = {
        .scheme = TPM2_ALG_OAEP,
        .details = {.oaep = {.hashAlg = TPM2_ALG_SHA256}}};
    TPM2B_DATA label = {0};

    // EncryptDecrypt Output Parameters
    TSS2L_SYS_AUTH_RESPONSE rsp_auths_array = {0};

    TSS2_RC rc =
        Tss2_Sys_RSA_Decrypt(sys_context, key_handle, &cmd_auths_array, in_data,
                             &in_scheme, &label, out_data, &rsp_auths_array);
    return rc;
}

TSS2_RC get_random_bytes(TSS2_SYS_CONTEXT *sys_ctx, uint8_t *buffer,
                         UINT16 size) {
    TPM2B_DIGEST random_bytes;
    TSS2_RC rc = Tss2_Sys_GetRandom(sys_ctx, NULL, size, &random_bytes, NULL);
    if (rc == TPM2_RC_SUCCESS) {
        memcpy(buffer, random_bytes.buffer, random_bytes.size);
    }
    return rc;
}

#define AES_KEY_SIZE 32
#define AES_IV_SIZE 16

// OpenSSL AES encryption
int ossl_aes_encrypt(const uint8_t *plaintext, int plen, const uint8_t *key,
                     const uint8_t *iv, uint8_t *ciphertext) {
    EVP_CIPHER_CTX *ctx;
    int len;

    // Create and initialize cipher context
    CHECK_OSSL((ctx = EVP_CIPHER_CTX_new()), "EVP_CIPHER_CTX_new");
    CHECK_OSSL(EVP_EncryptInit_ex(ctx, EVP_aes_256_cbc(), NULL, key, iv) == 1,
               "EVP_EncryptInit_ex");

    // Provide the message to be encrypted and obtain the encrypted output
    CHECK_OSSL(EVP_EncryptUpdate(ctx, ciphertext, &len, plaintext, plen) == 1,
               "EVP_EncryptUpdate");
    int clen = len;

    // Complete the encryption process
    CHECK_OSSL(EVP_EncryptFinal_ex(ctx, ciphertext + len, &len) == 1,
               "EVP_EncryptFinal_ex");
    clen += len;

    // Cleanup
    EVP_CIPHER_CTX_free(ctx);
    return clen;
}

// OpenSSL AES decryption
int ossl_aes_decrypt(const uint8_t *ciphertext, int clen, const uint8_t *key,
                     const uint8_t *iv, uint8_t *plaintext) {
    EVP_CIPHER_CTX *ctx;
    int len;

    // Create and initialize cipher context
    CHECK_OSSL((ctx = EVP_CIPHER_CTX_new()), "EVP_CIPHER_CTX_new");
    CHECK_OSSL(EVP_DecryptInit_ex(ctx, EVP_aes_256_cbc(), NULL, key, iv) == 1,
               "EVP_DecryptInit_ex");

    // Provide the message to be decrypted and obtain the decrypted output
    CHECK_OSSL(EVP_DecryptUpdate(ctx, plaintext, &len, ciphertext, clen) == 1,
               "EVP_DecryptUpdate");
    int plen = len;

    // Complete the decryption process
    CHECK_OSSL(EVP_DecryptFinal_ex(ctx, plaintext + len, &len) == 1,
               "EVP_DecryptFinal_ex");
    plen += len;

    // Cleanup
    EVP_CIPHER_CTX_free(ctx);
    return plen;
}

// Read file content into a buffer
size_t read_file_to_buffer(const char *filename, uint8_t **buffer) {
    FILE *file = fopen(filename, "rb");
    if (!file) {
        perror("Failed to open file");
        return 0;
    }

    fseek(file, 0, SEEK_END);
    size_t size = ftell(file);
    fseek(file, 0, SEEK_SET);

    *buffer = (uint8_t *)malloc(size);
    if (!*buffer) {
        fclose(file);
        perror("Failed to allocate memory");
        return 0;
    }

    fread(*buffer, 1, size, file);
    fclose(file);
    return size;
}

// Write buffer content to a file
size_t write_buffer_to_file(const char *filename, const uint8_t *buffer,
                            size_t size) {
    FILE *file = fopen(filename, "wb");
    if (!file) {
        perror("Failed to open file for writing");
        return 0;
    }

    size_t written = fwrite(buffer, 1, size, file);
    fclose(file);
    return written;
}

int main(int argc, char *argv[]) {

    TSS2_RC rc;

    // Initialize TCTI contexts
    rc = init_tcti_context(&tcti_context);
    CHECK_PASSED(rc);

    // Initialize SAPI context
    rc = init_sys_context(tcti_context, &sys_context);
    CHECK_PASSED(rc);

    // Startup TPM
    rc = startup_tpm(sys_context);
    CHECK_PASSED(rc);

    // ===================================
    // | Receiver: Generate RSA key pair |
    // ===================================
    const char *rsa_key_password = "HelloWorld";
    TPM2_HANDLE rsa_key_handle;
    TPM2B_PUBLIC rsa_public_key;
    rc = create_key_rsa_2048(sys_context, rsa_key_password, &rsa_key_handle,
                             &rsa_public_key);
    CHECK_PASSED(rc);
    handle_recorder.handles[handle_recorder.count++] =
        rsa_key_handle; // Record the RSA key handle
    printf("Receiver: Generated RSA key pair.\n");

    // ===================================
    // | Sender: Generate AES key and IV |
    // ===================================
    uint8_t aes_key[AES_KEY_SIZE];
    rc = get_random_bytes(sys_context, aes_key, AES_KEY_SIZE);
    CHECK_PASSED(rc);
    uint8_t aes_iv[AES_IV_SIZE];
    rc = get_random_bytes(sys_context, aes_iv, AES_IV_SIZE);
    CHECK_PASSED(rc);
    printf("Sender: Generated AES key and IV.\n");

    // ===========================================
    // | Sender: Encrypt AES key with RSA pubkey |
    // ===========================================
    TPM2B_PUBLIC_KEY_RSA aes_key_to_encrypt = {
        .size = AES_KEY_SIZE,
    };
    memcpy(aes_key_to_encrypt.buffer, aes_key, AES_KEY_SIZE);
    // rc = rsa_pubkey_encrypt(sys_context, NULL, rsa_key_handle_R,
    TPM2B_PUBLIC_KEY_RSA encrypted_aes_key_S = {0};
    rc = rsa_pubkey_encrypt(sys_context, NULL, rsa_key_handle,
                            &aes_key_to_encrypt, &encrypted_aes_key_S);
    CHECK_PASSED(rc);
    printf("Sender: Encrypted AES key with RSA public key.\n");

    // ======================================
    // | Sender: Encrypt data with RSA key  |
    // ======================================
    const char *pf = "红楼梦.txt";
    uint8_t *pbuf = NULL;
    size_t plen = read_file_to_buffer(pf, &pbuf);
    size_t csize = plen + EVP_MAX_BLOCK_LENGTH; // Ensure enough space for padding
    uint8_t* cbuf = (uint8_t *)malloc(csize);
    int clen = ossl_aes_encrypt(pbuf, plen, aes_key, aes_iv, cbuf);
    CHECK_OSSL(clen > 0, "ossl_aes_encrypt");
    printf("Sender: Encrypted data with AES.\n");
    write_buffer_to_file("红楼梦_encrypted.txt", cbuf, clen);

    // ==========================================
    // | Receiver: Decrypt AES key with RSA key |
    // ==========================================
    TPM2B_PUBLIC_KEY_RSA decrypted_aes_key = { .size = sizeof(TPM2B_PUBLIC_KEY_RSA) - 2, };
    rc = rsa_prvkey_decrypt(sys_context, rsa_key_password, rsa_key_handle,
                            &encrypted_aes_key_S, &decrypted_aes_key);
    CHECK_PASSED(rc);
    printf("Receiver: Decrypted AES key with RSA private key.\n");

    // =======================================
    // | Receiver: Decrypt data with AES key |
    // =======================================
    size_t dsize = csize;
    uint8_t* dbuf = (uint8_t *)malloc(dsize);
    int dlen = ossl_aes_decrypt(cbuf, clen, decrypted_aes_key.buffer,
                                aes_iv, dbuf);
    CHECK_OSSL(dlen > 0, "ossl_aes_decrypt");
    printf("Receiver: Decrypted data with AES key.\n");
    write_buffer_to_file("红楼梦_decrypted.txt", dbuf, dlen);

    // Cleanup resources
    cleanup(rc);
    free(pbuf);
    free(cbuf);
    free(dbuf);

    return 0;
}
